import torch
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image
from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
import os
import io
import base64
import cv2
import numpy as np
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
from sklearn.metrics import confusion_matrix, classification_report
import json
import warnings
import sys
from contextlib import redirect_stderr

# 抑制pytorch-grad-cam的AttributeError警告
warnings.filterwarnings('ignore')

# 创建一个包装类来抑制GradCAM的AttributeError
class SafeGradCAM(GradCAM):
    def __del__(self):
        try:
            super().__del__()
        except AttributeError:
            pass  # 忽略AttributeError

# 定义与训练时一致的模型架构
class TumorClassifier(nn.Module):
    def __init__(self, num_classes):
        super(TumorClassifier, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(32 * 56 * 56, 128),
            nn.ReLU(inplace=True),
            nn.Linear(128, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

# 初始化Flask应用
app = Flask(__name__, static_folder='client/build')
CORS(app)  # 允许跨域请求

# 全局变量存储模型
model = None
classes = ['glioma', 'meningioma', 'notumor', 'pituitary']

# Grad-CAM目标层（最后一个卷积层）
target_layers = None

# 定义数据预处理步骤（与训练时保持一致）
def get_transforms():
    return transforms.Compose([
        transforms.Resize((224, 224)),  # 调整图像大小
        transforms.ToTensor(),  # 转换为张量
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
    ])

# 加载模型
def load_model():
    global model, target_layers
    model = TumorClassifier(num_classes=4)
    
    # 加载训练好的权重
    try:
        model.load_state_dict(torch.load('best_model.pth', map_location=torch.device('cpu')))
        
        # 为Grad-CAM启用梯度计算
        for param in model.parameters():
            param.requires_grad = True
        
        # 设置Grad-CAM目标层（选择卷积层而不是激活层）
        # 找到最后一个卷积层
        conv_layers = []
        for module in model.features.modules():
            if isinstance(module, (torch.nn.Conv2d, torch.nn.ConvTranspose2d)):
                conv_layers.append(module)
        
        if conv_layers:
            target_layers = [conv_layers[-1]]  # 最后一个卷积层
        else:
            # 如果没有找到卷积层，使用第一个有参数的层
            for module in model.features.modules():
                if len(list(module.parameters())) > 0:
                    target_layers = [module]
                    break
        
        print(f"模型加载成功，目标层: {target_layers}")
    except Exception as e:
        print(f"无法加载模型权重: {str(e)}")
        return False
    
    return True

# 图像预处理函数
def preprocess_image(image_bytes):
    try:
        # 打开图像
        image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
        # 应用预处理
        transform = get_transforms()
        image_tensor = transform(image).unsqueeze(0)  # 添加批次维度
        return image_tensor
    except Exception as e:
        raise ValueError(f"图像预处理失败: {str(e)}")

# 预测函数
def predict_image(image_tensor):
    global model
    with torch.no_grad():
        outputs = model(image_tensor)
        probabilities = torch.nn.functional.softmax(outputs, dim=1)
        confidence, predicted = torch.max(probabilities, 1)
    
    # 获取预测结果
    predicted_class = classes[predicted.item()]
    confidence_score = confidence.item() * 100  # 转换为百分比
    
    # 获取所有类别的概率
    all_probabilities = []
    for i, prob in enumerate(probabilities[0]):
        all_probabilities.append({
            "class": classes[i],
            "confidence": round(prob.item() * 100, 2)
        })
    
    # 按置信度排序
    all_probabilities.sort(key=lambda x: x["confidence"], reverse=True)
    
    return {
        "predicted_class": predicted_class,
        "confidence": round(confidence_score, 2),
        "all_probabilities": all_probabilities
    }

# 生成Grad-CAM热力图
def generate_gradcam(image_tensor, predicted_class_idx, original_size=None):
    global model, target_layers
    
    if target_layers is None:
        print("目标层未设置")
        return None
    
    if model is None:
        print("模型未加载")
        return None
    
    cam = None
    try:
        # 确保输入张量启用梯度并设置正确的数据类型
        image_tensor = image_tensor.clone().detach().requires_grad_(True)
        
        # 确保模型在训练模式（为了梯度计算）
        model.train()
        
        # 创建SafeGradCAM对象
        cam = SafeGradCAM(model=model, target_layers=target_layers)
        
        # 生成热力图
        targets = [ClassifierOutputTarget(predicted_class_idx)]
        
        # 使用torch.enable_grad()包装以确保梯度计算
        with torch.enable_grad():
            grayscale_cam = cam(input_tensor=image_tensor, targets=targets)
        
        # 转换为PIL图像
        grayscale_cam = grayscale_cam[0, :]  # 移除批次维度
        heatmap = (grayscale_cam * 255).astype(np.uint8)
        
        # 转换为RGB格式
        heatmap_rgb = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
        
        # 转换为PIL图像
        heatmap_pil = Image.fromarray(cv2.cvtColor(heatmap_rgb, cv2.COLOR_BGR2RGB))
        
        # 如果提供了原始尺寸，调整到原始尺寸；否则保持224x224
        if original_size:
            heatmap_pil = heatmap_pil.resize(original_size, Image.LANCZOS)
        else:
            heatmap_pil = heatmap_pil.resize((224, 224), Image.LANCZOS)
        
        return heatmap_pil
        
    except Exception as e:
        print(f"Grad-CAM生成失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return None
    finally:
        # 恢复模型到评估模式
        if model is not None:
            model.eval()
        if cam is not None:
            del cam

# 图像转base64
def image_to_base64(image):
    buffer = io.BytesIO()
    image.save(buffer, format='PNG')
    img_str = base64.b64encode(buffer.getvalue()).decode()
    return f"data:image/png;base64,{img_str}"

# 批量预测函数
def predict_batch(image_tensors):
    global model
    results = []
    
    with torch.no_grad():
        for i, image_tensor in enumerate(image_tensors):
            try:
                outputs = model(image_tensor)
                probabilities = torch.nn.functional.softmax(outputs, dim=1)
                confidence, predicted = torch.max(probabilities, 1)
                
                predicted_class = classes[predicted.item()]
                confidence_score = confidence.item() * 100
                
                # 获取所有类别的概率
                all_probabilities = []
                for j, prob in enumerate(probabilities[0]):
                    all_probabilities.append({
                        "class": classes[j],
                        "confidence": round(prob.item() * 100, 2)
                    })
                
                all_probabilities.sort(key=lambda x: x["confidence"], reverse=True)
                
                results.append({
                    "index": i,
                    "predicted_class": predicted_class,
                    "confidence": round(confidence_score, 2),
                    "all_probabilities": all_probabilities
                })
            except Exception as e:
                results.append({
                    "index": i,
                    "error": f"预测失败: {str(e)}"
                })
    
    return results

# 批量预测函数（带热力图）
def predict_batch_with_heatmap(image_tensors, original_sizes=None):
    global model
    results = []
    
    with torch.no_grad():
        for i, image_tensor in enumerate(image_tensors):
            try:
                outputs = model(image_tensor)
                probabilities = torch.nn.functional.softmax(outputs, dim=1)
                confidence, predicted = torch.max(probabilities, 1)
                
                predicted_class = classes[predicted.item()]
                confidence_score = confidence.item() * 100
                
                # 获取所有类别的概率
                all_probabilities = []
                for j, prob in enumerate(probabilities[0]):
                    all_probabilities.append({
                        "class": classes[j],
                        "confidence": round(prob.item() * 100, 2)
                    })
                
                all_probabilities.sort(key=lambda x: x["confidence"], reverse=True)
                
                # 生成Grad-CAM热力图，传递原始尺寸
                original_size = original_sizes[i] if original_sizes and i < len(original_sizes) else None
                heatmap = generate_gradcam(image_tensor, predicted.item(), original_size)
                heatmap_base64 = image_to_base64(heatmap) if heatmap else None
                
                results.append({
                    "index": i,
                    "predicted_class": predicted_class,
                    "confidence": round(confidence_score, 2),
                    "all_probabilities": all_probabilities,
                    "heatmap": heatmap_base64
                })
            except Exception as e:
                results.append({
                    "index": i,
                    "error": f"预测失败: {str(e)}"
                })
    
    return results

# 健康检查接口
@app.route('/api/health', methods=['GET'])
def health():
    return jsonify({
        "status": "healthy",
        "model_loaded": model is not None
    })

# 预测接口
@app.route('/api/predict', methods=['POST'])
def predict():
    if model is None:
        return jsonify({"error": "模型未加载"}), 500
    
    if 'image' not in request.files:
        return jsonify({"error": "未提供图像文件"}), 400
    
    file = request.files['image']
    
    if file.filename == '':
        return jsonify({"error": "未选择图像文件"}), 400
    
    try:
        # 读取图像数据
        image_bytes = file.read()
        
        # 预处理图像
        image_tensor = preprocess_image(image_bytes)
        
        # 进行预测
        result = predict_image(image_tensor)
        
        return jsonify(result)
    except Exception as e:
        return jsonify({"error": f"预测过程中发生错误: {str(e)}"}), 500

# 带热力图的预测接口
@app.route('/api/predict_with_heatmap', methods=['POST'])
def predict_with_heatmap():
    if model is None:
        return jsonify({"error": "模型未加载"}), 500
    
    if 'image' not in request.files:
        return jsonify({"error": "未提供图像文件"}), 400
    
    file = request.files['image']
    
    if file.filename == '':
        return jsonify({"error": "未选择图像文件"}), 400
    
    try:
        # 读取图像数据
        image_bytes = file.read()
        
        # 获取原始图像尺寸
        original_image = Image.open(io.BytesIO(image_bytes))
        original_size = original_image.size  # (width, height)
        
        # 预处理图像
        image_tensor = preprocess_image(image_bytes)
        
        # 进行预测
        result = predict_image(image_tensor)
        
        # 生成热力图，传递原始尺寸
        predicted_class_idx = classes.index(result["predicted_class"])
        heatmap = generate_gradcam(image_tensor, predicted_class_idx, original_size)
        
        if heatmap:
            heatmap_base64 = image_to_base64(heatmap)
            result["heatmap"] = heatmap_base64
        else:
            result["heatmap"] = None
        
        return jsonify(result)
    except Exception as e:
        return jsonify({"error": f"预测过程中发生错误: {str(e)}"}), 500

# 批量预测接口
@app.route('/api/predict_batch', methods=['POST'])
def predict_batch_api():
    if model is None:
        return jsonify({"error": "模型未加载"}), 500
    
    if 'images' not in request.files:
        return jsonify({"error": "未提供图像文件"}), 400
    
    files = request.files.getlist('images')
    
    if not files or files[0].filename == '':
        return jsonify({"error": "未选择图像文件"}), 400
    
    try:
        # 预处理所有图像
        image_tensors = []
        for file in files:
            image_bytes = file.read()
            image_tensor = preprocess_image(image_bytes)
            image_tensors.append(image_tensor)
        
        # 批量预测
        results = predict_batch(image_tensors)
        
        # 统计各类别数量
        class_counts = {}
        for result in results:
            if "predicted_class" in result:
                class_name = result["predicted_class"]
                class_counts[class_name] = class_counts.get(class_name, 0) + 1
        
        return jsonify({
            "results": results,
            "class_counts": class_counts,
            "total_images": len(files)
        })
    except Exception as e:
        return jsonify({"error": f"批量预测过程中发生错误: {str(e)}"}), 500

# 批量预测接口（带热力图）
@app.route('/api/predict_batch_with_heatmap', methods=['POST'])
def predict_batch_with_heatmap_api():
    if model is None:
        return jsonify({"error": "模型未加载"}), 500
    
    if 'images' not in request.files:
        return jsonify({"error": "未提供图像文件"}), 400
    
    files = request.files.getlist('images')
    
    if not files or files[0].filename == '':
        return jsonify({"error": "未选择图像文件"}), 400
    
    try:
        # 预处理所有图像并保存原始尺寸
        image_tensors = []
        original_sizes = []
        for file in files:
            image_bytes = file.read()
            
            # 获取原始图像尺寸
            original_image = Image.open(io.BytesIO(image_bytes))
            original_sizes.append(original_image.size)
            
            image_tensor = preprocess_image(image_bytes)
            image_tensors.append(image_tensor)
        
        # 批量预测（带热力图）
        results = predict_batch_with_heatmap(image_tensors, original_sizes)
        
        # 统计各类别数量
        class_counts = {}
        for result in results:
            if "predicted_class" in result:
                class_name = result["predicted_class"]
                class_counts[class_name] = class_counts.get(class_name, 0) + 1
        
        return jsonify({
            "results": results,
            "class_counts": class_counts,
            "total_images": len(files)
        })
    except Exception as e:
        return jsonify({"error": f"批量预测过程中发生错误: {str(e)}"}), 500

# Grad-CAM功能展示接口
@app.route('/api/gradcam_demo', methods=['POST'])
def gradcam_demo():
    """
    Grad-CAM功能展示接口
    展示AI模型的可解释性，生成热力图显示模型关注的区域
    """
    if model is None:
        return jsonify({"error": "模型未加载"}), 500
    
    if 'image' not in request.files:
        return jsonify({"error": "未提供图像文件"}), 400
    
    file = request.files['image']
    
    if file.filename == '':
        return jsonify({"error": "未选择图像文件"}), 400
    
    try:
        # 读取图像数据
        image_bytes = file.read()
        
        # 预处理图像
        image_tensor = preprocess_image(image_bytes)
        
        # 进行预测
        prediction_result = predict_image(image_tensor)
        
        # 生成Grad-CAM热力图
        predicted_class_idx = classes.index(prediction_result["predicted_class"])
        heatmap = generate_gradcam(image_tensor, predicted_class_idx)
        
        # 将原始图像转换为base64
        original_image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
        original_image = original_image.resize((224, 224))
        original_base64 = image_to_base64(original_image)
        
        # 将热力图转换为base64
        heatmap_base64 = image_to_base64(heatmap) if heatmap else None
        
        # 创建叠加图像（原始图像 + 热力图）
        overlay_image = None
        if heatmap:
            # 将热力图转换为numpy数组
            heatmap_array = np.array(heatmap)
            original_array = np.array(original_image)
            
            # 创建叠加图像（可以调整透明度）
            alpha = 0.6  # 热力图透明度
            overlay = cv2.addWeighted(original_array, 1-alpha, heatmap_array, alpha, 0)
            overlay_pil = Image.fromarray(overlay)
            overlay_base64 = image_to_base64(overlay_pil)
        else:
            overlay_base64 = None
        
        # 返回详细的Grad-CAM展示结果
        result = {
            "success": True,
            "message": "Grad-CAM热力图生成成功",
            "prediction": prediction_result,
            "gradcam_info": {
                "target_class": prediction_result["predicted_class"],
                "target_class_index": predicted_class_idx,
                "confidence": prediction_result["confidence"],
                "explanation": f"模型预测为 {prediction_result['predicted_class']}，置信度 {prediction_result['confidence']:.2f}%"
            },
            "images": {
                "original": original_base64,
                "heatmap": heatmap_base64,
                "overlay": overlay_base64
            },
            "technical_details": {
                "method": "Grad-CAM (Gradient-weighted Class Activation Mapping)",
                "target_layer": "最后一个卷积层",
                "colormap": "JET (红色=高关注度，蓝色=低关注度)",
                "resolution": "224x224",
                "overlay_alpha": 0.6
            }
        }
        
        return jsonify(result)
        
    except Exception as e:
        return jsonify({
            "success": False,
            "error": f"Grad-CAM生成过程中发生错误: {str(e)}"
        }), 500

# Grad-CAM功能说明接口
@app.route('/api/gradcam_info', methods=['GET'])
def gradcam_info():
    """
    返回Grad-CAM功能的详细说明和技术信息
    """
    return jsonify({
        "title": "Grad-CAM (Gradient-weighted Class Activation Mapping)",
        "description": "AI模型可解释性技术，展示深度学习模型的决策过程",
        "features": [
            "实时生成热力图显示模型关注的区域",
            "红色区域表示高关注度，蓝色区域表示低关注度",
            "可调节透明度进行叠加显示",
            "帮助医生理解AI的判断依据",
            "提高AI诊断的可信度和透明度"
        ],
        "technical_details": {
            "algorithm": "Grad-CAM",
            "implementation": "pytorch-grad-cam",
            "target_layer": "最后一个卷积层",
            "colormap": "JET",
            "output_resolution": "224x224",
            "overlay_transparency": "可调节 (0.0-1.0)"
        },
        "use_cases": [
            "医学影像诊断辅助",
            "AI模型可解释性研究",
            "医生培训和教育",
            "模型性能验证",
            "诊断决策支持"
        ],
        "api_endpoints": {
            "demo": "/api/gradcam_demo",
            "predict_with_heatmap": "/api/predict_with_heatmap",
            "info": "/api/gradcam_info"
        }
    })

# 训练指标接口
@app.route('/api/training_metrics', methods=['GET'])
def training_metrics():
    try:
        # 读取真实训练数据
        training_results_path = os.path.join(os.path.dirname(__file__), '..', 'training_results', 'training_metrics.json')
        
        if os.path.exists(training_results_path):
            with open(training_results_path, 'r') as f:
                training_results = json.load(f)
            
            return jsonify(training_results)
        else:
            # 如果文件不存在，返回默认数据
            return jsonify({
                "error": "Training results not found. Please run training first.",
                "training_data": {
                    "epochs": [],
                    "train_loss": [],
                    "val_loss": [],
                    "train_acc": [],
                    "val_acc": []
                },
                "confusion_matrix": {
                    "labels": classes,
                    "matrix": []
                },
                "class_metrics": [],
                "model_info": {
                    "architecture": "Custom CNN",
                    "total_params": 0,
                    "input_size": "224x224x3",
                    "num_classes": 4,
                    "optimizer": "Adam",
                    "learning_rate": 0.001,
                    "batch_size": 16,
                    "epochs": 0
                }
            })
    except Exception as e:
        return jsonify({"error": f"Failed to load training metrics: {str(e)}"}), 500

# 为React应用提供服务
@app.route('/', defaults={'path': ''})
@app.route('/<path:path>')
def serve(path):
    if path != "" and os.path.exists(app.static_folder + '/' + path):
        return send_from_directory(app.static_folder, path)
    else:
        return send_from_directory(app.static_folder, 'index.html')

if __name__ == '__main__':
    # 加载模型
    if not load_model():
        print("模型加载失败，服务器将不提供预测功能")
    
    # 启动服务器
    app.run(host='0.0.0.0', port=5000, debug=True)